rm(list=ls())
############ initialize  ##########################
setwd("/Volumes/Disk2/Future_PM/Dataverse")
library(fields); library(maps); library(ncdf4);library(abind)
open.ncdf=nc_open;get.var.ncdf=ncvar_get;close.ncdf=nc_close
source('Function/get_geo.R')
source('Function/get_met.R')
source('Function/functions.R')

ss=load("Data/PM_region.Rdata")
time=seq(as.Date('1999-01-01'),as.Date('2013-12-31'),by='month')
date=cbind(as.numeric(format(time,"%Y")),as.numeric(format(time,"%m")))

#================= Part 1: read all data ==================
load("Data/monthly-prate_1999-2013.Rdata")#precip
	spdata=array(NA,dim(precip))
	for (month in 1:12){
		  ind=(date[,2]==month)
		  y= precip[,,ind]
		  for (i in 1:dim(y)[1]){
			  for (j in 1:dim(y)[2]){
				  spdata[i,j,ind]= mov.detrend(y[i,j,])
			  }
		  }
	  }	
precip=spdata

surf_T=read_met('air.mon.mean.nc')
surf_RH=read_met('rhum.mon.mean.nc')
surf_uwind=read_met('uwnd.mon.mean.nc')
surf_vwind=read_met('vwnd.mon.mean.nc')
T= surf_T$data
RH=surf_RH$data
u0=surf_uwind$data
v0=surf_vwind$data
lon.out= surf_T$lon-360
lat.out= surf_T$lat

load("Data/PM25_1999_2013_invdist_2.5x2.5_500km.Rdata")
total_PM[total_PM>=200]=200
total_PM[total_PM<=0]=0
for (year in 1999:2013){
	for (month in 1:12){
		ind=(total_time[,1]==year & total_time[,2]==month)
		num=apply(!is.na(total_PM[,,ind]),c(1,2),sum,na.rm=TRUE)
		avg=apply(total_PM[,,ind],c(1,2),mean,na.rm=TRUE)
		avg[num<=9]=NA
		if (year==1999 & month==1){
			monthly_PM=avg
		} else {
			monthly_PM=abind(monthly_PM,avg,along=3)
		}
	}
}

monthly_PM.std=array(NA,dim(monthly_PM))
for (month in 1:12){
	ind=(date[,2]==month)
	PM=monthly_PM[,,ind]
	for (i in 1:length(lon.frame)){
		for (j in 1:length(lat.frame)){
			monthly_PM.std[i,j,ind]= mov.detrend(PM[i,j,])
		}
	}
}
US_region= PM.region==1

#=== Part 2: select one grid box =======================
month=6
time.ind=(date[,2]>=(month-1) & date[,2]<=(month+1))
ii=20;jj=5
index= monthly_PM.std[ii,jj, time.ind]

lon0=-84.4
lat0=33.8
# loc=find.lon.lat(lon.frame[ii],lat.frame[jj],lon.out,lat.out)
loc=find.lon.lat(lon0,lat0,lon.out,lat.out)
ind1=seq(loc[1]-6,loc[1]+6)
ind2=seq(loc[2]-4,loc[2]+4)
corr1=find.cor2(T[ind1,ind2, time.ind],  index ,p_value=1.05)
corr2=find.cor2(RH[ind1,ind2, time.ind], index ,p_value=1.05)		
corr3=find.cor2(precip[ind1,ind2, time.ind], index ,p_value=1.05)
corr4=find.cor2(u0[ind1,ind2, time.ind], index ,p_value=1.05)
corr5=find.cor2(v0[ind1,ind2, time.ind], index ,p_value=1.05)

#######Figure 1###############
dev.new(width=6,height=5)
mai=c(0.1, 0.1, 0.2, 0.1)
mgp=c(1.2, 0.4, 0)
tcl=-0.2
ps=12

m <- rbind(c(0.1,0.9,0.8,1),
c(0,0.33,0.6,0.9), c(0.33, 0.67, 0.6, 0.9), c(0.67, 1, 0.6, 0.9),
				c(0.15, 0.48, 0.3, 0.6),c(0.48,0.81,0.3,0.6),
				c(0.2,0.8,0.23,0.53)
				)
close.screen(all.screens = TRUE)				
split.screen(m)
cr=find.Rlim(0.05, 45) 
screen(1)
par(mai=mai, mgp=mgp, tcl=tcl, ps=ps)
# text(x=0.5,y=0.6,"Correlation of PM2.5 with metorology")
text(x=0.5,y=0.6,bquote('Correlation of '*PM[2.5]~ 'with meteorology'))
screen(2)
par(mai=mai, mgp=mgp, tcl=tcl, ps=ps)
image(lon.out[ind1],lat.out[ind2], corr1,zlim=c(-0.75,0.75),col=rwb.colors(32),xaxt='n',yaxt='n',xlab='',ylab='');
map("world", add = TRUE)
points(lon.frame[ii],lat.frame[jj],cex=1.0)
title('(a) Temperature',cex.main=0.9,font.main=1)
sig=which(abs(corr1)>=cr,arr.ind=TRUE)
for(i in 1:dim(sig)[1]){
	points(lon.out[ind1][sig[i,1]],lat.out[ind2][sig[i,2]],pch=16,cex=0.3)	
}
box()
screen(3)
par(mai=mai, mgp=mgp, tcl=tcl, ps=ps)
image(lon.out[ind1],lat.out[ind2], corr2,zlim=c(-0.75,0.75),col=rwb.colors(32),xaxt='n',yaxt='n',xlab='',ylab='');
map("world", add = TRUE)
points(lon.frame[ii],lat.frame[jj],cex=1.0)
title('(b) Relative humidity',cex.main=0.9,font.main=1)
sig=which(abs(corr2)>=cr,arr.ind=TRUE)
for(i in 1:dim(sig)[1]){
	points(lon.out[ind1][sig[i,1]],lat.out[ind2][sig[i,2]],pch=16,cex=0.3)	
}
box()
screen(4)
par(mai=mai, mgp=mgp, tcl=tcl, ps=ps)
image(lon.out[ind1],lat.out[ind2], corr3,zlim=c(-0.75,0.75),col=rwb.colors(32),xaxt='n',yaxt='n',xlab='',ylab='');
map("world", add = TRUE)
points(lon.frame[ii],lat.frame[jj],cex=1.0)
title('(c) Precipitation',cex.main=0.9,font.main=1)
sig=which(abs(corr3)>=cr,arr.ind=TRUE)
for(i in 1:dim(sig)[1]){
	points(lon.out[ind1][sig[i,1]],lat.out[ind2][sig[i,2]],pch=16,cex=0.3)	
}
box()
screen(5)
par(mai=mai, mgp=mgp, tcl=tcl, ps=ps)
image(lon.out[ind1],lat.out[ind2], corr4,zlim=c(-0.75,0.75),col=rwb.colors(32),xaxt='n',yaxt='n',xlab='',ylab='');
map("world", add = TRUE)
points(lon.frame[ii],lat.frame[jj],cex=1.0)
title('(d) East-west wind',cex.main=0.9,font.main=1)
sig=which(abs(corr4)>=cr,arr.ind=TRUE)
for(i in 1:dim(sig)[1]){
	points(lon.out[ind1][sig[i,1]],lat.out[ind2][sig[i,2]],pch=16,cex=0.3)	
}
box()
screen(6)
par(mai=mai, mgp=mgp, tcl=tcl, ps=ps)
image(lon.out[ind1],lat.out[ind2], corr5,zlim=c(-0.75,0.75),col=rwb.colors(32),xaxt='n',yaxt='n',xlab='',ylab='');
map("world", add = TRUE)
points(lon.frame[ii],lat.frame[jj],cex=1.0)
title('(e) North-south wind',cex.main=0.9,font.main=1)
sig=which(abs(corr5)>=cr,arr.ind=TRUE)
for(i in 1:dim(sig)[1]){
	points(lon.out[ind1][sig[i,1]],lat.out[ind2][sig[i,2]],pch=16,cex=0.3)	
}
box()
screen(7)
 mai=c(0.1, 0.1, 0.05, 0.2)
 par(mai=mai, mgp=mgp, tcl=tcl, ps=9)
image.plot(corr5,legend.shrink=0.9,legend.only=TRUE,zlim=c(-0.75,0.75),col=rwb.colors(32),horizontal=TRUE,legend.width=2.5,axis.args = list(mgp = c(0, 0.5, 0),tcl=-0.2,padj=-1
))
text("r",x=par("usr")[2]-0.22,y=par("usr")[3]+0.26,srt=0, adj = 0,xpd=TRUE,cex=1.8)

  

#=========== Figure 2 ================================
F=abind(corr1,corr2,corr3,corr4,corr5, along=3);F[is.na(F)]=0
size=dim(F)
new.F<- matrix(F, nrow=size[1]*size[2], ncol=size[3], byrow=FALSE)
ss=svd(new.F)
uu=ss$u   #spacial loading  
vv=ss$v   #variable loading
dd=ss$d   #eigen value
ap=uu%*%diag(dd)%*%t(vv)
uu[,2]=-uu[,2]
vv[,2]=-vv[,2]

dev.new(width=6.5,height=4)
par(mar=c(1,3,1,3))
par(mfrow=c(2,2))
for (i in 1:2){
sp=matrix(uu[,i],nrow=size[1],ncol=size[2],byrow=FALSE)
plot.field(sp,lon.out[ind1],lat.out[ind2],type='sign',legend.width=2,zlim=c(-0.3,0.3), mai=c(0.1, 0.3,0.2,0.2))
box()
if(i==1){
	text(x=-72,y=25,c('32%'),font=1,cex=1.3)
    mtext("SVD1",side=2,col=1,line=0.3,cex=1, font=1) 	
    title('Spatial weight',font.main=1)
    text(x=-99,y=45,c('(a)'),font=1,cex=1.2)
} 
if(i==2){
	text(x=-72,y=25,c('31%'),font=1,cex=1.3)
    mtext("SVD2",side=2,col=1,line=0.3,cex=1, font=1) 		
    text(x=-99,y=45,c('(c)'),font=1,cex=1.2)
} 
points(lon.frame[ii],lat.frame[jj],pch=16,col=1)

mai=c(0.3, 0.6, 0.2, 0.2)
mgp=c(1.8, 0.4, 0)
par(mai=mai,mgp=mgp)
plot(vv[,i],type='h',xlab='',ylab='loadings',lwd=8,col=4,xlim=c(0.8, 5.2), ylim=c(-1,1),axes=FALSE);
box()
axis(2,col="black",las=1,lwd=1,tck=-0.01,cex.axis=1.0)  ## las=1 makes horizontal labels
axis(1,c(1,2,3,4,5),tck=-0.03,cex.main=1.0,labels=FALSE)
text(y=-1.3,x=c(1,2,3,4,5),labels=c('T','RH','precip','EW','NS'),xpd=TRUE,srt=0,cex.main=1.2,cex=1)
abline(h=0)
if(i==1) {
	title('Variable weight',font.main=1)
	text(x=0.95,y=0.95,c('(b)'),font=1,cex=1.2)
	text(x=4.9,y=-0.95,c('SVD1'),font=1,cex=1.2)	
}
if(i==2) {
	text(x=0.95,y=0.95,c('(d)'),font=1,cex=1.2)
	text(x=4.9,y=-0.95,c('SVD2'),font=1,cex=1.2)	
}
}
